use std::collections::HashSet;
use std::error::Error;
use std::fmt::Debug;

use async_trait::async_trait;
use axum::http;
use axum::http::request::Parts;
use axum::response::IntoResponse;

use crate::auth::permission::PermissionType;

pub trait AuthUser: Clone + Send + Sync + Debug {
    fn get_permission(&self) -> HashSet<String>;
}

pub trait AuthEndpoint: Clone + Send + Sync {
    fn auth_fail(&self, parts: &mut Parts) -> impl IntoResponse;
    fn permission_fail(&self, parts: &mut Parts,
                       permission_type: PermissionType,
                       permissions: HashSet<String>) -> impl IntoResponse {
        self.auth_fail(parts)
    }
    fn auth_error<E: Error>(&self, parts: &mut Parts, error: E) -> impl IntoResponse {
        self.auth_fail(parts)
    }
}

#[derive(Debug, Clone, Default)]
pub struct DefaultAuthEndpoint;

impl AuthEndpoint for DefaultAuthEndpoint {
    fn auth_fail(&self, parts: &mut Parts) -> impl IntoResponse {
        http::StatusCode::UNAUTHORIZED
    }
}

#[async_trait]
pub trait AuthnBackend: Clone + Send + Sync {
    type User: Send + Sync + Clone + Debug;
    type Error: Error + Send + Sync;

    async fn authenticate(&self, parts: &mut Parts) -> Result<Option<Self::User>, Self::Error>;

    async fn get_permission(&self, user: &Self::User) -> HashSet<String>;
}


#[derive(Debug, Clone)]
pub struct AuthSession<Backend: AuthnBackend> {
    pub user: Option<Backend::User>,
    pub backend: Backend,
    pub permission_type: PermissionType,
    permissions: Option<HashSet<String>>,
}

impl<Backend: AuthnBackend> AuthSession<Backend> {
    async fn get_permission(&mut self) -> Option<&HashSet<String>> {
        match &self.user {
            None => {
                return None;
            }
            Some(user) => {
                if self.permissions.is_none() {
                    let permissions = self.backend.get_permission(user).await;
                    self.permissions = Some(permissions);
                    self.permissions.as_ref()
                } else {
                    return self.permissions.as_ref();
                }
            }
        }
    }
}

pub mod middleware {
    use std::collections::HashSet;
    use std::future::Future;
    use std::pin::Pin;
    use std::task::{Context, Poll};

    use axum::http::Request;
    use axum::response::IntoResponse;
    use tower::Service;

    use crate::auth::{AuthEndpoint, AuthnBackend, AuthSession, DefaultAuthEndpoint};
    use crate::auth::permission::FilterMode;

    #[derive(Clone)]
    pub struct AuthManagerLayer<BK, AE> where BK: AuthnBackend, AE: AuthEndpoint {
        filter_mode: FilterMode,
        backend: BK,
        auth_endpoint: AE,
    }

    impl<BK, AE> AuthManagerLayer<BK, AE> where BK: AuthnBackend, AE: AuthEndpoint {
        pub fn default_config(backend: BK) -> AuthManagerLayer<BK, DefaultAuthEndpoint> {
            AuthManagerLayer {
                filter_mode: FilterMode::default(),
                backend,
                auth_endpoint: DefaultAuthEndpoint::default(),
            }
        }
        pub fn new(filter_mode: FilterMode, backend: BK, auth_endpoint: AE) -> Self {
            Self {
                filter_mode,
                backend,
                auth_endpoint,
            }
        }
    }

    #[derive(Clone, Debug)]
    pub struct AuthManager<T, BK, AE> where BK: AuthnBackend, AE: AuthEndpoint {
        inner: T,
        filter_mode: FilterMode,
        backend: BK,
        auth_endpoint: AE,
    }

    impl<T, ReqBody, BK, AE> Service<Request<ReqBody>> for AuthManager<T, BK, AE>
        where
            T: Service<Request<ReqBody>> + Clone + Send + 'static,
            T::Future: Send + 'static,
            T::Response: IntoResponse + Send,
            ReqBody: Send + 'static,
            BK: AuthnBackend + 'static,
            AE: AuthEndpoint + 'static,
    {
        type Response = axum::response::Response;
        type Error = T::Error;
        type Future = Pin<Box<dyn Future<Output=Result<Self::Response, Self::Error>> + Send>>;

        fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            self.inner.poll_ready(cx)
        }

        fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
            let clone = self.inner.clone();
            let backend = self.backend.clone();
            let auth_endpoint = self.auth_endpoint.clone();
            let permission_type = self.filter_mode.match_url(req.uri().path()).clone();
            // take the service that was ready
            let mut inner = std::mem::replace(&mut self.inner, clone);
            Box::pin(
                async move {
                    let mut req = req;
                    let mut session_user: Option<BK::User> = None;
                    let mut session_permissions: Option<HashSet<String>> = None;
                    if !permission_type.skip_auth() {
                        let (mut parts, body) = req.into_parts();
                        match backend.authenticate(&mut parts).await {
                            Ok(None) => {
                                return Ok(auth_endpoint.auth_fail(&mut parts).into_response());
                            }
                            Ok(Some(user)) => {
                                if !permission_type.skip_permission() {
                                    let permissions = backend.get_permission(&user).await;
                                    if !permission_type.match_permission(&permissions) {
                                        // 失败: 权限不足
                                        return Ok(auth_endpoint.permission_fail(&mut parts, permission_type, permissions).into_response());
                                    }
                                    session_permissions = Some(permissions);
                                }
                                session_user = Some(user);
                            }
                            Err(e) => {
                                //     失败: 鉴权异常
                                return Ok(auth_endpoint.auth_error(&mut parts, e).into_response());
                            }
                        };
                        req = Request::from_parts(parts, body);
                    }

                    req.extensions_mut().insert(AuthSession {
                        user: session_user,
                        backend,
                        permission_type,
                        permissions: session_permissions,
                    });
                    inner.call(req).await.map(|rep| {
                        rep.into_response()
                    })
                }
            )
        }
    }

    impl<T, BK, AE> tower::Layer<T> for AuthManagerLayer<BK, AE> where BK: AuthnBackend, AE: AuthEndpoint {
        type Service = AuthManager<T, BK, AE>;

        fn layer(&self, service: T) -> Self::Service {
            AuthManager {
                inner: service,
                filter_mode: self.filter_mode.clone(),
                backend: self.backend.clone(),
                auth_endpoint: self.auth_endpoint.clone(),
            }
        }
    }
}

pub mod permission {
    use std::collections::{HashMap, HashSet};

    #[derive(Clone, Debug)]
    pub enum PermissionType {
        // 不需要登录
        Anonymous,
        // 不需要权限,登录就可以
        Authenticated,
        // 任意一个
        AnyPermission(HashSet<String>),
        // 完全匹配
        HasPermission(String),
        // 完全不匹配
        NotPermission(String),
        // 必须都匹配
        AllPermission(HashSet<String>),
    }

    impl PermissionType {
        pub fn skip_auth(&self) -> bool {
            matches!(self, PermissionType::Anonymous)
        }
        pub fn skip_permission(&self) -> bool {
            matches!(self, PermissionType::Authenticated)
        }

        // 权限匹配
        pub fn match_permission(&self, user_permissions: &HashSet<String>) -> bool {
            match self {
                PermissionType::Anonymous | PermissionType::Authenticated => {
                    true
                }
                PermissionType::AnyPermission(permissions) => {
                    !permissions.is_disjoint(permissions)
                }
                PermissionType::HasPermission(permission) => {
                    user_permissions.contains(permission)
                }
                PermissionType::NotPermission(permission) => {
                    !user_permissions.contains(permission)
                }
                PermissionType::AllPermission(permissions) => {
                    permissions.is_subset(user_permissions)
                }
            }
        }
    }

    #[derive(Clone, Debug)]
    pub struct FilterMode {
        pub default_mode: PermissionType,
        pub route_table: Option<HashMap<String, PermissionType>>,
    }

    impl Default for FilterMode {
        fn default() -> Self {
            FilterMode {
                default_mode: PermissionType::Authenticated,
                route_table: None,
            }
        }
    }

    impl FilterMode {
        pub fn with_default_mode(mut self, default_mode: PermissionType) -> Self {
            self.default_mode = default_mode;
            self
        }
        pub fn add_route(mut self, url: impl Into<String>, mode: PermissionType) -> Self {
            self.route_table.get_or_insert_with(HashMap::new).insert(url.into(), mode);
            self
        }

        pub fn match_url(&self, url: &str) -> &PermissionType {
            self.route_table.as_ref()
                .and_then(|table| table.get(url))
                .unwrap_or(&self.default_mode)
        }
    }
}


pub mod jwt {
    use std::collections::HashSet;

    use async_trait::async_trait;
    use axum::http;
    use axum::http::HeaderName;
    use axum::http::request::Parts;
    use jsonwebtoken::{DecodingKey, Validation};
    use serde::de::DeserializeOwned;
    use tracing::info;

    use crate::auth::{AuthnBackend, AuthUser};

    #[derive(Clone)]
    pub struct JwtBackend<U: AuthUser + DeserializeOwned> {
        token_header_name: HeaderName,
        decode_key: DecodingKey,
        validation: Validation,
        user_type: std::marker::PhantomData<U>,
    }

    impl<U: AuthUser + DeserializeOwned> JwtBackend<U> {
        pub fn new(
            token_header_name: HeaderName,
            secret: &[u8],
            validation: Validation,
        ) -> Self {
            Self {
                token_header_name,
                decode_key: DecodingKey::from_secret(secret),
                validation,
                user_type: std::marker::PhantomData,
            }
        }

        pub fn new_by_key(
            secret: &[u8]
        ) -> Self {
            Self {
                token_header_name: http::header::AUTHORIZATION,
                decode_key: DecodingKey::from_secret(secret),
                validation: Validation::default(),
                user_type: std::marker::PhantomData,
            }
        }
    }

    #[async_trait]
    impl<U> AuthnBackend for JwtBackend<U> where U: AuthUser + DeserializeOwned
    {
        type User = U;
        type Error = std::convert::Infallible;

        async fn authenticate(&self, parts: &mut Parts) -> Result<Option<Self::User>, Self::Error> {
            info!("header_name: {},value: {:?}", self.token_header_name, parts.headers.get(&self.token_header_name));
            Ok(
                parts.headers.get(&self.token_header_name)
                    .and_then(|token_value| token_value.to_str().ok())
                    .and_then(|token| {
                        let r = jsonwebtoken::decode::<U>(token, &self.decode_key, &self.validation);
                        info!("token: {:?}", r);
                        r.ok()
                    })
                    .map(|token_data| token_data.claims)
            )
        }


        async fn get_permission(&self, user: &U) -> HashSet<String> {
            user.get_permission()
        }
    }
}
